#
# requires pycrypto [http://www.amk.ca/python/code/crypto.html]
#
try:
    from Crypto.PublicKey.RSA import generate as _rsa_generate  
    from Crypto.PublicKey.RSA import construct as _rsa_construct  
except ImportError:
    def _rsa_generate(*args):
        raise ImportError("Crypto.PublicKey.RSA not available")
    _rsa_construct = _rsa_generate

_randpool = None
def get_random_bytes(count):
    global _randpool
    if not _randpool:
        from Crypto.Util.randpool import RandomPool 
        _randpool = RandomPool()
    return _randpool.get_bytes(count)

def _encode_number(num):
    bytes = ""
    while num > 0:
        bytes += chr(num & 0xff)
        num >>= 8
    return bytes[::-1].encode("base64").replace("\n", "")
def _decode_number(base64):
    bytes = base64.decode("base64")
    num = 0
    for i, b in enumerate(bytes):
        num <<= 8
        num += ord(b)
    return long(num)
def _encode_numbers(numbers):
    return ":".join(_encode_number(n) for n in numbers)
def _decode_numbers(raw):
    return tuple(_decode_number(b64) for b64 in raw.split(":"))

class RSAKey(object):
    __slots__ = ["_key"]
    RSA_BITS = 2048
    PUB_KEY_ATTRS = ["n", "e"]
    PRI_KEY_ATTRS = ["n", "e", "d", "p", "q", "u"]
    
    def __init__(self, _key):
        self._key = _key
    def __repr__(self):
        return "<RSAKey %s, %s>" % self._get_public_numbers()
    def __hash__(self):
        return hash(self._get_public_numbers())
    def __cmp__(self, other):
        return cmp(self._get_public_numbers(), other._get_public_numbers())
    @classmethod
    def _from_numbers(cls, nums):
        lpub = len(cls.PUB_KEY_ATTRS)
        lpri = len(cls.PRI_KEY_ATTRS)
        if len(nums) in (lpub, lpri):
            return cls(_rsa_construct(nums))
        else:
            raise ValueError("expected %d or %d numbers, %d given" % (lpub, lpri, len(nums)))
    @classmethod
    def generate(cls, bits = RSA_BITS):
        return cls(_rsa_generate(bits, get_random_bytes))
    @classmethod
    def load(cls, data):
        return cls._from_numbers(_decode_numbers(data))
    
    def _get_public_numbers(self):
        return tuple(getattr(self._key, k) for k in self.PUB_KEY_ATTRS)
    def _get_private_numbers(self):
        return tuple(getattr(self._key, k) for k in self.PRI_KEY_ATTRS)
    def dump_public(self):
        return _encode_numbers(self._get_public_numbers())
    def dump_private(self):
        return _encode_numbers(self._get_private_numbers())
    
    def get_public_key(self):
        return self._from_numbers(self._get_public_numbers())
    def sign(self, data):
        return _encode_number(self._key.sign(data, "")[0])
    def verify(self, data, sig):
        return self._key.verify(data, [_decode_number(sig)])


#
# to prevent timing attacks, create a generic key so that verifying with 
# BadKey and AnyKey will consume the same amount of time as with a real key
#
#_some_key = None
#def _get_some_rsa_key():
#    global _some_key
#    if not _some_key:
#        _some_key = RSAKey.generate()
#    return _some_key

class BadRSAKey(object):
    def verify(self, data, sig):
        #_get_some_rsa_key().verify(data, sig)
        return False
BadRSAKey = BadRSAKey()

class AnyRSAKey(object):
    def verify(self, data, sig):
        #_get_some_rsa_key().verify(data, sig)
        return True
AnyRSAKey = AnyRSAKey()



